#include <iostream>
#include <string>
using namespace std;

#include "imcrypto.h"
#include "openssl/evp.h"

static inline EVP_CIPHER_CTX* _getContext(void* pData)
{
	return (EVP_CIPHER_CTX*)pData;
}

static inline const EVP_CIPHER* _getCipher(int ct, int pt)
{
	switch(ct)
	{
	case IMCrypto::CT_NULL: return EVP_enc_null();
	case IMCrypto::CT_DES:
	{
		switch(pt)
		{
		case IMCrypto::PT_ECB:	return EVP_des_ecb();
		case IMCrypto::PT_CFB64:return EVP_des_cfb64();
		case IMCrypto::PT_CFB1:	return EVP_des_cfb1();
		case IMCrypto::PT_CFB8:	return EVP_des_cfb8();
		case IMCrypto::PT_OFB:	return EVP_des_ofb();
		case IMCrypto::PT_CBC:	return EVP_des_cbc();
		}
		return NULL;
	}
	case IMCrypto::CT_DES3:
	{
		switch(pt)
		{
		case IMCrypto::PT_NULL:	return EVP_des_ede3();
		case IMCrypto::PT_ECB:	return EVP_des_ede3_ecb();
		case IMCrypto::PT_CFB64:return EVP_des_ede3_cfb64();
		case IMCrypto::PT_CFB1:	return EVP_des_ede3_cfb1();
		case IMCrypto::PT_CFB8:	return EVP_des_ede3_cfb8();
		case IMCrypto::PT_OFB:	return EVP_des_ede3_ofb();
		case IMCrypto::PT_CBC:	return EVP_des_ede3_cbc();
		}
		return NULL;
	}
	case IMCrypto::CT_AES128:
	{
		switch(pt)
		{
		case IMCrypto::PT_ECB:	return EVP_aes_128_ecb();
		case IMCrypto::PT_CFB1:	return EVP_aes_128_cfb1();
		case IMCrypto::PT_CFB8:	return EVP_aes_128_cfb8();
		case IMCrypto::PT_CFB128:return EVP_aes_128_cfb128();
		case IMCrypto::PT_OFB:	return EVP_aes_128_ofb();
		case IMCrypto::PT_CBC:	return EVP_aes_128_cbc();
		}
		return NULL;
	}
	case IMCrypto::CT_AES192:
	{
		switch(pt)
		{
		case IMCrypto::PT_ECB:	return EVP_aes_192_ecb();
		case IMCrypto::PT_CFB1:	return EVP_aes_192_cfb1();
		case IMCrypto::PT_CFB8:	return EVP_aes_192_cfb8();
		case IMCrypto::PT_CFB128:return EVP_aes_192_cfb128();
		case IMCrypto::PT_OFB:	return EVP_aes_192_ofb();
		case IMCrypto::PT_CBC:	return EVP_aes_192_cbc();
		}
		return NULL;
	}
	case IMCrypto::CT_AES256:
	{
		switch(pt)
		{
		case IMCrypto::PT_ECB:	return EVP_aes_256_ecb();
		case IMCrypto::PT_CFB1:	return EVP_aes_256_cfb1();
		case IMCrypto::PT_CFB8:	return EVP_aes_256_cfb8();
		case IMCrypto::PT_CFB128:return EVP_aes_256_cfb128();
		case IMCrypto::PT_OFB:	return EVP_aes_256_ofb();
		case IMCrypto::PT_CBC:	return EVP_aes_256_cbc();
		}
		return NULL;
	}
	default:
		return NULL;
	}

	return NULL;
}

IMCrypto::IMCrypto() : m_ctx(NULL)
{
}

IMCrypto::~IMCrypto()
{
	if ( m_ctx )
	{
		EVP_CIPHER_CTX* c((EVP_CIPHER_CTX*)m_ctx);
		EVP_CIPHER_CTX_cleanup(c);
		delete c;
		m_ctx = NULL;
	}
}

bool
IMCrypto::s_getSize(size_t* keysize, size_t* ivsize, size_t* blocksize, CRYPTO_TYPE ct, PADDING_TYPE pt)
{
	const EVP_CIPHER* cipher = _getCipher(ct,pt);
	if ( NULL == cipher ) return false;

	if ( keysize ) *keysize = EVP_CIPHER_key_length(cipher);
	if ( ivsize ) *ivsize = EVP_CIPHER_iv_length(cipher);
	if ( blocksize ) *blocksize = EVP_CIPHER_block_size(cipher);

	return true;
}

IMCrypto*
IMCrypto::s_createEncrypt(CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	const EVP_CIPHER* cipher(_getCipher(ct, pt));
	if ( cipher == NULL ) return NULL;

	EVP_CIPHER_CTX* ctx(new EVP_CIPHER_CTX);
	if ( ctx == NULL ) return NULL;

	do
	{
		EVP_CIPHER_CTX_init(ctx);
		if ( padding_size != (size_t)-1 ) EVP_CIPHER_CTX_set_padding(ctx, padding_size);

		if ( 0 == EVP_EncryptInit_ex(ctx, cipher, NULL, key, iv) ) break;

		IMCrypto* crypto(new IMCrypto);
		if ( NULL == crypto ) break;

		crypto->m_ctx = ctx;
		return crypto;
	} while (false);

	EVP_CIPHER_CTX_cleanup(ctx);
	delete ctx;
	return NULL;
}

IMCrypto*
IMCrypto::s_createDecrypt(CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	const EVP_CIPHER* cipher(_getCipher(ct, pt));
	if ( cipher == NULL ) return NULL;

	EVP_CIPHER_CTX* ctx(new EVP_CIPHER_CTX);
	if ( ctx == NULL ) return NULL;

	do
	{
		EVP_CIPHER_CTX_init(ctx);

		if ( 0 == EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv) ) break;

		if ( padding_size != (size_t)-1 ) EVP_CIPHER_CTX_set_padding(ctx, padding_size);

		IMCrypto* crypto(new IMCrypto);
		if ( NULL == crypto ) break;

		crypto->m_ctx = ctx;
		return crypto;
	} while (false);

	EVP_CIPHER_CTX_cleanup(ctx);
	delete ctx;
	return NULL;
}

void
IMCrypto::s_destroy(IMCrypto* obj)
{
	delete obj;
}

size_t
IMCrypto::getKeySize(void) const
{
	return (size_t)EVP_CIPHER_CTX_key_length(_getContext(m_ctx));
}

size_t
IMCrypto::getIVSize(void) const
{
	return (size_t)EVP_CIPHER_CTX_iv_length(_getContext(m_ctx));
}

size_t
IMCrypto::getBlockSize(void) const
{
	return (size_t)EVP_CIPHER_CTX_block_size(_getContext(m_ctx));
}

bool
IMCrypto::encrypt(std::string& out, const char* buf, size_t blen) const
{
	char* obuf = (char*)malloc(((blen/getBlockSize())+1)*getBlockSize());
	if ( NULL == obuf ) return false;

	size_t olen;
	if ( !encrypt(obuf, olen, buf, blen) )
	{
		free(obuf);
		return false;
	}

	if ( olen > 0 ) out.assign(obuf, olen);
	else out.clear();

	free(obuf);
	return true;
}

bool
IMCrypto::encrypt(char* out, size_t& _olen, const char* buf, size_t blen) const
{
	unsigned char* obuf((unsigned char*)out);

	int olen(0);
	if ( 0 == EVP_EncryptUpdate(_getContext(m_ctx), obuf, &olen, (const unsigned char*)buf, blen) ) return false;

	int tlen(0);
	if ( 0 == EVP_EncryptFinal_ex(_getContext(m_ctx), obuf+olen, &tlen) ) return false;

	_olen = olen + tlen;

	return true;
}

bool
IMCrypto::decrypt(std::string& out, const char* buf, size_t blen) const
{
	char* obuf = (char*)malloc(blen);
	if ( NULL == obuf ) return false;

	size_t olen;
	if ( !decrypt(obuf, olen, buf, blen) )
	{
		free(obuf);
		return false;
	}

	if ( olen > 0 ) out.assign(obuf, olen);
	else out.clear();

	free(obuf);
	return true;
}

bool
IMCrypto::decrypt(char* out, size_t& _olen, const char* buf, size_t blen) const
{
	unsigned char* obuf((unsigned char*)out);

	int olen(0);
	if ( 0 == EVP_DecryptUpdate(_getContext(m_ctx), obuf, &olen, (const unsigned char*)buf, blen) ) return false;

	int tlen(0);
	if ( 0 == EVP_DecryptFinal_ex(_getContext(m_ctx), obuf+olen, &tlen) ) return false;

	_olen = olen + tlen;

	return true;
}

bool
IMCrypto::s_encrypt(char* obuf, size_t& olen, const char* buf, size_t blen, CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	IMCrypto* crypto(IMCrypto::s_createEncrypt(ct,pt,key,iv,padding_size));
	if ( NULL == crypto ) return false;

	bool res(crypto->encrypt(obuf, olen, buf, blen) );
	s_destroy(crypto);
	return res;
}

bool
IMCrypto::s_encrypt(std::string& out, const char* buf, size_t blen, CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	IMCrypto* crypto(IMCrypto::s_createEncrypt(ct,pt,key,iv,padding_size));
	if ( NULL == crypto ) return false;

	bool res(crypto->encrypt(out, buf, blen) );
	s_destroy(crypto);
	return res;
}

bool
IMCrypto::s_decrypt(char* obuf, size_t& olen, const char* buf, size_t blen, CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	IMCrypto* crypto(IMCrypto::s_createDecrypt(ct,pt,key,iv,padding_size));
	if ( NULL == crypto ) return false;

	bool res(crypto->decrypt(obuf, olen, buf, blen) );
	s_destroy(crypto);
	return res;
}

bool
IMCrypto::s_decrypt(std::string& out, const char* buf, size_t blen, CRYPTO_TYPE ct, PADDING_TYPE pt, const unsigned char* key, const unsigned char* iv, size_t padding_size)
{
	IMCrypto* crypto(IMCrypto::s_createDecrypt(ct,pt,key,iv,padding_size));
	if ( NULL == crypto ) return false;

	bool res(crypto->decrypt(out, buf, blen) );
	s_destroy(crypto);
	return res;
}

